from transformers import BertTokenizer

#加载字典和分词工具
token = BertTokenizer.from_pretrained(
    r"../../my_model_cache/bert-base-chinese/models--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f")
print(token)

# 分词单个序列
inputs = token.tokenize("Hello, world!")
print(inputs)
print("==============================")
# 分词序列对
# 使用 __call__ 方法批量处理字符串列表
inputs = token(["Hello, world!", "How are you?"])
print(inputs)
print("==============================")
sents = ["白日依山尽",
         "这本书没读到底，不是特别喜欢。完全可以用序中的评价来表达我的感受：可以包容，却不想实践。除了物质，我们要追求的还有很多。"]

#批量编码句子
out = token.batch_encode_plus(
    batch_text_or_text_pairs=[sents[0], sents[1]],
    add_special_tokens=True,
    # 当句子长度大于max_lenght时，截断
    truncation=True,
    max_length=50,
    # 一律补0到max_length长度
    padding="max_length",
    # 可取值为tf,pt,np,默认为返回list
    return_tensors=None,
    return_attention_mask=True,
    return_token_type_ids=True,
    return_special_tokens_mask=True,
    # 返回length长度
    return_length=True
)
#input_ids 就是编码后的词
#token_type_ids第一个句子和特殊符号的位置是0，第二个句子的位置1
#special_tokens_mask 特殊符号的位置是1，其他位置是0
for k, v in out.items():
    print(k, ":", v)

print("===================================")
#解码文本数据
print(token.decode(out["input_ids"][0]))
print(token.decode(out["input_ids"][1]))
